# -*- coding: utf-8 -*-
"""
Created on Tue Nov 23 18:44:43 2021

@author: zhiyinan
"""

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import Parameters_1D as par_RNN 
from Polar1D_vectorized import *
import van_Genuchten1D_vectorized as vg 
from scipy import integrate
from sklearn.preprocessing import MinMaxScaler
from casadi import *
import matplotlib.font_manager as font_manager
import time
import pandas as pd
from sklearn.linear_model import Ridge
import random

# %% Define NN layers
def sigmoid(x):
    return 1/(1+np.exp(-x))

def LSTMlayer_sig(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    # _c = np.tanh(s_t[:,2*hunit:3*hunit])
    _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    # h_t = o*np.tanh(c_t)
    h_t = o*sigmoid(c_t)
    return h_t, c_t

def LSTMlayer_tanh(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    _c = np.tanh(s_t[:,2*hunit:3*hunit])
    # _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    h_t = o*np.tanh(c_t)
    # h_t = o*sigmoid(c_t)
    return h_t, c_t




# %% Define sub model 2
model_2 = tf.keras.models.load_model("Modeling/p20_f1_10s_tanh_2h_021_032_8_90_e8_n")

def M2(x):
    weightLSTM_1 = model_2.layers[0].get_weights()
    warr1, uarr1, barr1 = weightLSTM_1

    weightD = model_2.layers[1].get_weights()
    kernel2, bias2 = weightD
    
    n1 = 10
    # Initialise the model with the hidden states
    h_t1 = SX.zeros((1, n1))
    # Long term memory
    C_t1 = SX.zeros((1, n1))
    for i in range(x.shape[0]):
        xi = x[i,:].reshape((1,2))
        L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
        h_t1 = vertcat(h_t1, L1[0])
        C_t1 = vertcat(C_t1, L1[1])
    
    #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
    # Dense layer
    out2 = []
    for i in range(bias2.shape[0]):
        temp2 = mtimes(h_t1[-1,:], kernel2)[i]
    temp2 = bias2.reshape((1, bias2.shape[0])) + mtimes(h_t1[-1,:], kernel2)
    out2 = tanh(temp2)     
    
    return out2

# %% Input Signal generation
np.random.seed(2)
random.seed(2)


soilPars = loamySoil() # The soil parameters
totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints = par_RNN.spatialVariables_1D()  #Spatial parameters
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = \
                                 par_RNN.temporalVariable_dataGen(30000*2*60*60) #time parameters

# lbu = -5.10e-05  #  unit [m/s]
# ubu = -4.94e-05  #  unit [m/s]

# irr_freq = int(20*60*60/samplingTime)   # PRBS change once only if a minimum of 100 hr is reached

uv  = np.zeros((totTimeSteps,1))


i=0
while i <= 30000:
    irr_freq = int(20*60*60/samplingTime)
    ulist = np.array([-1.2e-6, -1e-6, -0.4e-6, -2.5e-6, -4.5e-6, -1.8e-6, -0.6e-6, -1.4e-6])/5
    # ulist = [-2/1000/3600, -1.5/1000/3600, -0.9/1000/3600, -1.2/1000/3600, -0.5/1000/3600]
    t = irr_freq + random.randrange(irr_freq) # (1+random.randrange(1))*
    uv[i:i+t] = random.choice(ulist)
    i = i+t
    print(i)
    

plt.figure(1)
plt.plot(uv)                


initCondition = -0.2*np.ones(totalNodes)
#Arrays to store the results




def ODE(t, head, irrigAmount, cropCoeff, refEvap, soilPars):
    return Richards1D_vectorized(head, irrigAmount, cropCoeff, refEvap, soilPars)

cropCoeff = 0.88*np.ones(totTimeSteps)
refEvap = 3.102e-08*np.ones(totTimeSteps)

# %% Integration 

headArray= np.zeros((totTimeSteps+1, totalNodes)) # Pressure head, the state
headArray[0,:] = initCondition

volMoisture = np.zeros((totTimeSteps+1, totalNodes)) # Volumetric moisture, the measured output
volMoisture[0,:] = volMoistureAllNodes_1D(initCondition, soilPars).ravel()


i=1
temp = headArray[i-1]
ui = uv[i-1]
for t in range(int(samplingTime/samplingTimeInternal)): 
    sol = integrate.solve_ivp(ODE, args = (ui, cropCoeff[i-1], refEvap[i-1], soilPars),
                              t_span=[0,samplingTimeInternal],y0=tuple(temp),method='BDF')
    temp = sol.y[:,-1]
    # ui = 0
headArray[i,:]=sol.y[:,-1]
volMoisture[i,:] = volMoistureAllNodes_1D(headArray[i,:], soilPars).ravel()

for i in range(2, totTimeSteps + 1):  # totTimeSteps + 1
    print(i)
    temp = headArray[i-1]
    ui = uv[i-1]
    for t in range(int(samplingTime/samplingTimeInternal)): 
        sol = integrate.solve_ivp(ODE, args = (ui, cropCoeff[i-1], refEvap[i-1], soilPars),
                              t_span=[0,samplingTimeInternal],y0=tuple(temp),method='BDF')
        temp = sol.y[:,-1]
    # print(sol.y[:,-1])
        # ui = 0
    headArray[i,:]=sol.y[:,-1]
    volMoisture[i,:] = volMoistureAllNodes_1D(headArray[i,:], soilPars).ravel()

plt.figure(2)
plt.plot(volMoisture[:1000,19])

# np.savetxt("u_20hrIF_30000_021_032.txt", uv)
# np.savetxt("x_20hrIF_30000_021_032.txt", headArray)
# np.savetxt("y_20hrIF_30000_021_032.txt", volMoisture)

# %%
b = 1000
uv = np.loadtxt("Modeling/u_vali_20hrIF_1000_021_032.txt")
yv = np.loadtxt("Modeling/y_vali_20hrIF_1000_021_032.txt")[1:b+1,19]

# y_act = volMoisture[1:b+1,19]
yn = 0.2*(max(yv) - min(yv))*np.random.choice([-1,1],(yv.shape))*np.random.random((yv.shape))
yv_n = yv*(1+yn)

y_test = yv
x_test = np.zeros((b,2))
x_test[:,0] = uv[:b]
x_test[:,1] = yv_n[:b]


scaler_xt = MinMaxScaler()
x_test = scaler_xt.fit_transform(x_test)
y_test = y_test.reshape((y_test.shape[0],1))
scaler_yt = MinMaxScaler()
y_test = scaler_yt.fit_transform(y_test)

past = 20
future = 1

def multivariate_data(dataset, target, start_index, end_index, history_size,
                      target_size,u_shape): #, single_step=False):
    data = []
    labels = []

    start_index = start_index + history_size
    if end_index is None:
        end_index = len(dataset) - target_size
        
    for i in range(start_index, end_index):
        temp_p = dataset[i-history_size:i,:]
        for j in range(target_size-1):
        # temp_p = np.concatenate(((dataset[i-history_size:i,:]).flatten(), 
        #                            (dataset[i:i+target_size-1,:u_shape]).flatten()))
            temp = np.array([dataset[i+j,0], dataset[i,1]]).reshape((1,2))
            temp_p = np.concatenate((temp_p,temp))
        data.append(temp_p)
        labels.append((target[i:i+target_size].flatten()))
    

    return np.array(data), np.array(labels)

x_t, y_t = multivariate_data(x_test, y_test, 0, None, past, future,1)

# %%
y_pt_1 = model_2.predict(x_t)
ypt_1 = scaler_yt.inverse_transform(y_pt_1)
y_act = scaler_yt.inverse_transform(y_t.reshape((y_t.shape[0], y_t.shape[1])))

csfont = {'fontname':'Times New Roman'}
font = font_manager.FontProperties(family='Times New Roman',
                                    # weight='bold',
                                    style='normal', size=15)
plt.figure(3)
plt.plot(y_act[:,0], label = 'Actual Data')
plt.plot(ypt_1[:,0], label = 'Single-step-ahead Prediction')
# plt.plot(ypt_1[-500:,-1], label = 'Multi-step-ahead NN')
# plt.plot(ypt_2[30:,0], label = 'Multi-step-ahead NN -1st')
# plt.plot(ypt_2[:,-1], label = 'Multi-step-ahead NN - 30th')
plt.xlabel('Time Steps', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)
plt.legend(frameon=False, ncol=1, prop=font)

# %% Multi-step ahead prediction performance

t_pred = 20  # 30 sampling times 
error = np.zeros(x_t.shape[0])
y_save = np.zeros((a,t_pred))

for t0 in range(a):
    print(t0)
    y_sim_1 = np.zeros((t_pred,1))
    NN0_1 = x_t[t0,:]
    # y_sub = np.array([DM(M2(NN0_1)), DM(M2(NN0_1)), DM(M2(NN0_1))])  #, DM(M4(NN0_1)), DM(M5(NN0_1))])
    y_sim_1[0,:] = np.array(DM(M2(NN0_1))).flatten()
    u_sim = x_t[t0:t0+t_pred,-1,0]
    for i in range(t_pred-1):
        NN0_1 = np.concatenate((NN0_1[1:], np.array([u_sim[i], y_sim_1[i,0]]).reshape((1,2))))
        # y_sub = np.array([DM(M2(NN0_1)), DM(M2(NN0_1)), DM(M2(NN0_1))])  #, DM(M4(NN0_1)), DM(M5(NN0_1))])
        y_sim_1[i+1,:] = np.array(DM(M2(NN0_1))).flatten()
    yp_1 = scaler_yt.inverse_transform(y_sim_1).flatten()
    y_save[t0,:] = yp_1
    error[t0] = sum((y_act[t0:t0+t_pred].flatten() - yp_1)**2/y_act[t0:t0+t_pred].flatten())
# yp_2 = scaler_yt.inverse_transform(y_sim_2)
# yp = (yp_1 + yp_2)/2
# %% Plotting

font = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=17)
tplot = np.arange(a)
b = 190
c = 203

fig,ax1 = plt.subplots()
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
ax1.plot(tplot, y_act[:a], label = 'Actual Trajectory', color = "k", linewidth = 3)
ax1.plot(tplot, y_save[:,0], label = "Single-step-ahead") # 0.000914381408118862
# ax1.plot(tplot[:], y_save[:,4], label = "5-step-ahead", color = "r")  # 0.0017256467027993415
ax1.plot(tplot[9:], y_save[:a-9,9], label = "Ten-step-ahead", color = "r", linestyle = "dashdot") # 0.002444138910936656
ax1.plot(tplot[19:], y_save[:a-19,19], label = "Twenty-step-ahead", color = "y", linestyle = "dashed") # 0.0030556959486465352
ax1.legend(frameon=False, bbox_to_anchor=(0.1,1,1,0), ncol=2, prop=font)
plt.xlabel("Time Steps", **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)
plt.show()

# ax2 = plt.axes([0,0,1,1])
# ip = InsetPosition(ax1, [0.23,0.78,0.15,0.2])
# ax2.set_axes_locator(ip)
# mark_inset(ax1, ax2, loc1=2, loc2=4, fc="none", ec='0.6')
# ax2.plot(tplot[b:c+20], y_act[b:c+20], label = 'Actual Trajectory', color = "k", linewidth = 3)
# ax2.plot(tplot[b:c], y_save[b:c,0], label = "Single-step-ahead") 
# ax2.plot(tplot[b:c+9], y_save[b-9:c,9], label = "Ten-step-ahead", color = "r", linestyle = "dashdot") 
# ax2.plot(tplot[b:c+19], y_save[b-19:c,19], label = "Twenty-step-ahead", color = "y", linestyle = "dashed") 
# # ax2.legend(loc=0)
# # ax2.set_xticklabels(ax2.get_xticks(), backgroundcolor='w')
# ax2.tick_params(axis='x', which='major', pad=8)
# plt.suptitle('Composite Plot with original subsystems(Ca)', fontsize=15)



# y_save = np.loadtxt("Modeling/y_pred_multi_20hrIF_1000_021_032.txt")
# font = font_manager.FontProperties(family='Times New Roman',
#                                    # weight='bold',
#                                     style='normal', size=17)
# tplot = np.arange(a)
# plt.figure(4)
# plt.plot(tplot, y_act[:a], label = 'Actual Trajectory', color = "k", linewidth = 4)
# plt.plot(tplot, y_save[:,0], label = "1-step-ahead") # 0.0002940232003526744
# # plt.plot(tplot[:], y_save[:,4], label = "5-step-ahead", color = "r")  # 0.005604632140300995
# # plt.plot(tplot[:], y_save[:,9], label = "10-step-ahead", color = "y") # 0.0063427657634302585
# plt.plot(tplot[:], y_save[:,19], label = "20-step-ahead", color = "g") # 0.006432134933479153
# plt.legend(frameon=False, bbox_to_anchor=(0.3,0.2,0.3,0.3), ncol=1, prop=font)
# plt.xlabel("Time Steps", **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)

e_0 = np.sqrt(sum((y_act[:900,0] - y_save[:,0])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0])) #(max((y_act[:900,:] - y_save[:,:]).flatten()))
e_5 = np.sqrt(sum((y_act[4:904,0] - y_save[:,4])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0])) 
e_10 = np.sqrt(sum((y_act[9:909,0] - y_save[:,9])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0]))
e_20 = np.sqrt(sum((y_act[19:919,0] - y_save[:,19])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0]))

print(e_0*100)
print(e_5*100)
print(e_10*100)
print(e_20*100)

# e_0 = np.sqrt(sum((y_act[:a,0] - y_save[:,0])**2/y_act[:a,0]**2))/a
# e_5 = np.sqrt(sum((y_act[4:a+4,0] - y_save[:,4])**2/y_act[4:a+4,0]**2))/a
# e_10 = np.sqrt(sum((y_act[9:a+9,0] - y_save[:,9])**2/y_act[9:a+9,0]**2))/a
# e_20 = np.sqrt(sum((y_act[19:a+19,0] - y_save[:,19])**2/y_act[19:a+19,0]**2))/a

# np.savetxt("u_vali_20hrIF_1000_021_032.txt", uv)
# np.savetxt("x_vali_20hrIF_1000_021_032.txt", headArray)
# np.savetxt("y_vali_20hrIF_1000_021_032.txt", volMoisture)
# np.savetxt("y_pred_multi_20hrIF_1000_021_032.txt", y_save)


# %%
u = np.loadtxt("Modeling/u_20hrIF_30000_021_032.txt")
y = np.loadtxt("Modeling/y_20hrIF_30000_021_032.txt")[:,y_index]
yn = 0.1*(max(y) - min(y))*np.random.choice([-1,1],(y.shape))*np.random.random((y.shape))
y_n = y*(1+yn)

fig, axs = plt.subplots(2,sharex = True, figsize=(11,9))
a = 200
b = 800
t = np.arange(b-a+1)
axs[0].step(t[:b-a],-u[a:b], color = "k") 
# axs[1].plot(t[:a],y[:a]) #, label = 'Product Concentration')
axs[0].set_ylabel('Irrigation Rate ($m/s$)', fontsize='xx-large', labelpad=20, fontproperties=font)
axs[1].set_xlabel("Time Steps", fontsize='xx-large', fontproperties=font)
axs[1].set_ylabel('Soil Moisture ($m^3/m^3$)', fontsize='xx-large', labelpad=5, fontproperties=font)
# for ax in axs.flatten():
#     labels = ax.get_xticklabels() + ax.get_yticklabels()
#     [label.set_fontname('Times New Roman') for label in labels]
#     [label.set_fontsize('xx-large') for label in labels]
    

csfont = {'fontname':'Times New Roman'}
# plt.figure(2,figsize=(6, 4.5))
# plt.style.use('grayscale')
axs[1].plot(t, y[a:b+1], label = "Nominal", color = "k",linewidth = 3)
axs[1].plot(t, y_n[a:b+1], label = "10% White Noise", color = "r")
# plt.xlabel('Time (min)', **csfont, fontsize = 17)
# plt.ylabel('Product concentration ($mole/m^3$)', **csfont, fontsize = 17)
font = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=17)
axs[1].legend(frameon=False, ncol=1, loc="upper right", prop=font)
plt.show()

